# -*- coding: utf-8 -*-
"""
Created on Thu Nov 23 15:06:44 2023

update green function with machine.xlsx


@author: xiangguchn
"""

import os
import pandas as pd
import json
import numpy as np
import math

import multiprocessing as mp
from scipy import special
from scipy import interpolate
import matplotlib.pyplot as plt

# import settings
# from utils.xlsx_tools import get_part
from xlsx_tools import get_part


class green_function():
    """
    Green function is used to calculate magnetic flux with plasma. It is mutual induction
    matrix between different part in tokamak including plasma, pf coil, VV and limiter
    
    refer: 
        1. Lao 1977
        2. 两圆线圈间互感及耦合系数讨论
    
    Developed by xiangguchn 2023/12/14
    If you have some question, don't hesitate to ask xiangguchn@163.com for help.
    """
    
    
    
    def __init__(self):

        # self.machine_file = settings.MACHINE_FILE
        # self.sheet_name = settings.MACHINE_NAME
        self.sheet_name = 'EXL50U'
        self.machine_file = 'C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\machine.xlsx'
        # self.machine_file = 'C:\sepy\\data\\EXL50U\\machine.xlsx'
        self.cmu = np.pi*2.0e-7;
        self.num_cores = int(mp.cpu_count())
        self.isplot = 1 # plot to show the result of green function

    
    def get_machine_parameters(self):
        """ 
        used to get machine parameters from machine.xlsx file
        path ='C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\machine.xlsx'
        
        input: 
            option name         data_type    size    unit    meaning
            1      machine_name String
        
        output:
            option  name        data_type   size    unit    meaning
            1       machine     struct      1       a.u     struct with machine's parameters
        
        reference: 
        
        Developed by xiangguchn 2023/11/23
        Modified by qidongkai 2023/12/14
        If you have some question, don't hesitate to ask xiangguchn@163.com for help.
        """

        machine = {}
        df = pd.read_excel(self.machine_file, sheet_name=self.sheet_name, header=None)
        # local_info = settings.MACHINE_INFO[self.sheet_name]
    
        # 
        local_info = {'plasma':'A32', 'PF':'A1', 'TF':'A108', 'TFP':'A89', 'VV':'A142', 'limiter':'A70' }
        
    
        # plasma parameters        
        plasma_read = get_part(df, local_info['plasma'])

        plasma = {}
        for _, row in plasma_read.iterrows():
            plasma[row['plasma parameter']] = row['value']
        machine['plasma'] = plasma
        
        # PF coils
        PF_read = get_part(df, local_info['PF'])
        machine['PF'] = PF_read
        
        # TF         
        TF_read = get_part(df, local_info['TF'])
        machine['TF'] = TF_read
        
        # TF parameters          
        TFP_read = get_part(df, local_info['TFP'])
        TFP = {}
        for _, row in TFP_read.iterrows():
            TFP[row['name']] = row['value']
        machine['TFP'] = TFP
        
        
        # VV
        VV_read = get_part(df, local_info['VV'])
        machine['VV'] = VV_read
        
        
        # limiter
        limiter_read = get_part(df, local_info['limiter'])
        
        x_limiter, y_limiter = self.interp_limiter_location(limiter_read, 
                                                            plasma['topLimiter'])
        machine['limiter'] = {}
        machine['limiter']['x_limiter'] = x_limiter
        machine['limiter']['y_limiter'] = y_limiter
        machine['limiter']['limiter_read'] = limiter_read
        
        self.machine = machine
        
        return machine
    
    
    def interp_limiter_location(self, limiter_read, topLimiter):
        """ 
        limiter coordinates are not uniform in machine.xlsx file. 
        interpration is need to calculate mutual induction matrix
                
        input: 
            option name         data_type    size    unit    meaning
            1      limiter_read struct                       limiter position of X and Y
            1      topLimiter   float        1       m       limiter limits, not all limiter is used
        
        output:
            option  name        data_type   size    unit    meaning
            1       x_limiter   array       1*n     m       limiter's position in x direction
            1       y_limiter   array       1*n     m       limiter's position in y direction
        
        reference: 
        
        Developed by xiangguchn 2023/11/23
        If you have some question, don't hesitate to ask xiangguchn@163.com for help.
        """
        
        x = limiter_read['X']
        y = limiter_read['Y']
        
        y_interp = np.arange(-topLimiter, topLimiter, 0.01)
        
        x_mid = (max(x)+min(x))/2
        
        
        x_in = x[x<x_mid]
        y_in = y[x<x_mid]
        
        f = interpolate.interp1d(y_in, x_in)        
        x_in_interp = f(y_interp)
        
        
        
        x_out = x[x>x_mid]
        y_out = y[x>x_mid]
        f = interpolate.interp1d(y_out, x_out)        
        x_out_interp = f(y_interp)
        
        x_limiter = np.concatenate((x_in_interp, x_out_interp))
        y_limiter = np.concatenate((y_interp,    y_interp))
        
        return x_limiter, y_limiter
    
    
    def update_green_function(self):
        """
        calculate green funciton(mutual induction) between different 
        
        case number of mutual inductance between:
                  plasma  PF  VV  limiter
        plasma    1       2   3   7
        PF                4   5
        VV                    6
        limiter   8       9
        
        input: 
            option name         data_type    size    unit    meaning
            1      limiter_read struct                       limiter position of X and Y
            1      topLimiter   float        1       m       limiter limits, not all limiter is used
        
        output:
            option  name        data_type   size    unit    meaning
            1       x_limiter   array       1*n     m       limiter's position in x direction
            1       y_limiter   array       1*n     m       limiter's position in y direction
        
        reference: 
        
        Developed by xiangguchn 2023/11/23
        If you have some question, don't hesitate to ask xiangguchn@163.com for help.
        """
        
        # set
        eq = {}
        eq['green'] = {}
        eq['green']['grid'] = {}
        eq['green']['delta'] = 1
        eq['green']['Lsol']  = 1
        eq['green']['alpha'] = 1
        
        # get machine data from machine.xlsx
        machine = self.get_machine_parameters()
        eq['machine'] = machine
        
        
        # position of plasma
        num_plasma = machine['plasma']['nr'] * machine['plasma']['nz']
        x_plasma = np.linspace(machine['plasma']['XStart'], machine['plasma']['XEnd'], machine['plasma']['nr'])
        y_plasma = np.linspace(machine['plasma']['YStart'], machine['plasma']['YEnd'], machine['plasma']['nz'])
    
        [x_plasma, y_plasma] = np.meshgrid(x_plasma, y_plasma)
        x_plasma = np.reshape(x_plasma,  (num_plasma))
        y_plasma = np.reshape(y_plasma,  (num_plasma))
        
        
        eq['green']['grid']['plasma'] = {}        
        eq['green']['grid']['plasma']['x'] = x_plasma
        eq['green']['grid']['plasma']['y'] = y_plasma
    
    
        # grid gap, 0.23 is optimized coef for statistic results from Song Xianming
        gapx = (x_plasma[1]-x_plasma[0])*0.23
    
        
    
        # case 1: mutual inductance between plasma and plasma
        # if ~isequal(g_fit.machine.plamsa, machine.plasma):
        if 1:
            
            # position of source field
            x2 = x_plasma;
            y2 = y_plasma;
        
            # calculate mutual inductance between plasma coils    
            mi_plasma_plasma = np.zeros(shape=(num_plasma, num_plasma));
            
            pool = mp.Pool(self.num_cores)            
            mi_plasma_plasma = [pool.apply(self.calculate_mi_plasma_plasma, 
                                           args=(x_plasma, y_plasma, x2[i], y2[i], gapx)) 
                                for i in range(num_plasma)]
            
            eq['green']['mi_plasma_plasma'] = np.array(mi_plasma_plasma)    
            
            if self.isplot:
                plt.figure()
                plt.contour(np.reshape(x_plasma, (129, 129)), 
                            np.reshape(y_plasma, (129, 129)), 
                            np.reshape(sum(mi_plasma_plasma), (129, 129)))
                plt.title('mi_plasma_plasma')
                ax = plt.gca()
                ax.set_aspect(1)
            
            print('1: mi_plasma_plasma is ok')
        
        # case 2: mutual inductance between plasma and pf coils
        # if ~isequal(g_fit.machine.PF, machine.PF):
        if 1:
            
            num_pf = len(machine['PF'])
            
            # source field if pf coils
        
            # calculate mutual inductance between plasma coils    
            mi_plasma_pf = np.zeros(shape=(num_plasma, num_plasma));
            
            
            # for index in range(num_pf):
            #     mi = self.calculate_mi_plasma_pf(x_plasma, y_plasma, index, gapx)
            #     print(len(mi))
            
            pool = mp.Pool(self.num_cores)
            
            mi_plasma_pf = [pool.apply(self.calculate_mi_plasma_pf, 
                                            args=(x_plasma, y_plasma, index, gapx)) 
                                for index in range(num_pf)]
            
            eq['green']['mi_plasma_pf'] = np.array(mi_plasma_pf)    
            
            
            if self.isplot:
                plt.figure()
                plt.contour(np.reshape(x_plasma, (129, 129)), 
                            np.reshape(y_plasma, (129, 129)), 
                            np.reshape(sum(mi_plasma_pf), (129, 129)))
                plt.title('mi_plasma_pf')
                ax = plt.gca()
                ax.set_aspect(1)
            
            print('2: mi_plasma_pf is ok')
        
        
        # case 7: mutual inductance between plasma and limiter
        if 1:
            
            x_limiter = machine['limiter']['x_limiter']
            y_limiter = machine['limiter']['y_limiter']
            num_limiter = len(machine['limiter']['x_limiter'])
            
            # source field if pf coils
        
            # calculate mutual inductance between plasma coils    
            mi_plasma_limiter = np.zeros(shape=(num_plasma, num_limiter));
            
            
            # for index in range(num_pf):
            #     mi = self.calculate_mi_plasma_pf(x_plasma, y_plasma, index, gapx)
            #     print(len(mi))
            
            pool = mp.Pool(self.num_cores)
            
            mi_plasma_limiter = [pool.apply(self.calculate_mi_plasma_limiter, 
                                            args=(x_plasma, y_plasma, x_limiter[i], y_limiter[i], gapx)) 
                                for i in range(num_limiter)]
            
            eq['green']['mi_plasma_limiter'] = np.array(mi_plasma_limiter)    
            
            
            if self.isplot:
                plt.figure()
                plt.contour(np.reshape(x_plasma, (129, 129)), 
                            np.reshape(y_plasma, (129, 129)), 
                            np.reshape(sum(mi_plasma_limiter), (129, 129)))
                plt.title('mi_plasma_limiter')
                
                ax = plt.gca()
                ax.set_aspect(1)
            
            print('7: mi_plasma_limiter is ok')
        
        
        # case 8: mutual inductance between limiter and plasma
        if 1:
            
            x_limiter = machine['limiter']['x_limiter']
            y_limiter = machine['limiter']['y_limiter']
            num_limiter = len(machine['limiter']['x_limiter'])
            
            # source field if pf coils
        
            # calculate mutual inductance between plasma coils    
            mi_limiter_plasma = np.zeros(shape=(num_limiter, num_plasma));
            # for i in range(num_plasma):
            #     print(i)
            #     mi_limiter_plasma[:,i] = self.calculate_mi_limiter_plasma(x_limiter, y_limiter, x_plasma[i], y_plasma[i], gapx)
            
            
            pool = mp.Pool(self.num_cores)            
            mi_limiter_plasma = [pool.apply(self.calculate_mi_limiter_plasma, 
                                            args=(x_limiter, y_limiter, x_plasma[i], y_plasma[i], gapx)) 
                                for i in range(num_plasma)]
                        
            eq['green']['mi_limiter_plasma'] = np.array(mi_limiter_plasma)    
            
            # if self.isplot:
            #     plt.figure()
            #     plt.contour(np.reshape(x_limiter, (129, 129)), 
            #                 np.reshape(y_limiter, (129, 129)), 
            #                 np.reshape(sum(mi_limiter_plasma), (129, 129)))
            #     plt.title('mi_plasma_limiter')
                       
            
            print('8: mi_limiter_plasma is ok')
        
        
        # case 9: mutual inductance between limiter and pf
        if 1:
            
            x_limiter = machine['limiter']['x_limiter']
            y_limiter = machine['limiter']['y_limiter']
            num_limiter = len(machine['limiter']['x_limiter'])
            
            # source field if pf coils
        
            # calculate mutual inductance between plasma coils    
            mi_limiter_pf = np.zeros(shape=(num_limiter, num_plasma));
            
            
            # for index in range(num_pf):
            #     mi = self.calculate_mi_plasma_pf(x_plasma, y_plasma, index, gapx)
            #     print(len(mi))
            
            pool = mp.Pool(self.num_cores)
            
            mi_limiter_pf = [pool.apply(self.calculate_mi_limiter_pf, 
                                            args=(x_limiter, y_limiter, index, gapx)) 
                                for index in range(num_pf)]
            
            eq['green']['mi_limiter_pf'] = np.array(mi_limiter_pf)    
            print('9: mi_limiter_pf is ok')
            
            
            # if isplot:
            #     plt.figure()
            #     plt.contour()
            
            # mi_limiter_pf = [pool.apply(self.calculate_mi_limiter_pf, 
            #                                 args=(x_limiter, y_limiter, x_pf[i], y_pf[i], gapx)) 
            #                     for i in range(num_plasma)]            
            # eq['green']['mi_limiter_plasma'] = np.array(mi_limiter_plasma)    
        
        # save result to json
        # eq_json = json.dumps(eq)
        # f2 = open('C:\SEPY\\data\\EXL50U\\eq.json', 'w')
        # f2.write(eq_json)
        # f2.close()
        
        # np.save('C:\\sepy\\data\\EXL50U\\green_function.npy', 
        #         eq, allow_pickle=True)
        # np.save(settings.GREEN_FILE, eq, allow_pickle=True)
        # np.save('C:\\sepy\\data\\EXL50U\\green_function.npy', 
        #         eq, allow_pickle=True)
        np.save('C:\\Users\\xiang\\Documents\\XINAO\\AI\\SEPY\\data\\EXL50U\\green_function.npy', 
                eq, allow_pickle=True)
        
        return eq
        
    
    
    def calculate_mi_limiter_plasma(self, x_limiter, y_limiter, x_plasma, y_plasma, gapx):
        """ 
        calculate mutual induction between limiter and plasma
                
        input: 
            option name         data_type    size    unit    meaning
            1       x_limiter   array        1*n     m       limiter's position in x direction
            1       y_limiter   array        1*n     m       limiter's position in y direction
            1       x_plasma    array        1       m       plasma's position in x direction
            1       y_plasma    array        1       m       plasma's position in y direction
            1       gap         float        1       m       gap for same coil or 
        
        output:
            option  name                data_type   size    unit    meaning
            1       mi_limiter_plasma   array       1*n     m       mutual induction between limiter and plasma
        
        reference: 
        
        Developed by xiangguchn 2023/11/23
        If you have some question, don't hesitate to ask xiangguchn@163.com for help.
        """

        R1 = np.sqrt(((x_limiter+x_plasma)**2 + (y_limiter-y_plasma)**2).astype('float'))
        m = 4 * x_plasma * x_limiter / R1**2

        # avoiding the overlay of source and field
        index = np.where(abs(m-1)<1e-10)
    
        
        xx1 = x_limiter; # may modify, not exist outside [if ~isempty(index)]
        if index is not None:
            xx1[index] = x_limiter[index] + gapx;    # move 1 mm outside
            R1 = np.sqrt(((xx1+x_plasma)**2 + (y_limiter-y_plasma)**2).astype('float'))
            m  = 4 * x_plasma * xx1 / R1**2;
        
        
        myk = special.ellipk(m.astype('float'))
        mye = special.ellipe(m.astype('float'))
        mi_limiter_plasma = self.cmu * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_limiter_plasma
    
    
    def calculate_mi_plasma_limiter(self, x_plasma, y_plasma, x2, y2, gapx):


        R1 = np.sqrt((x_plasma+x2)**2 + (y_plasma-y2)**2);
        m = 4 * x2 * x_plasma / R1**2;

        # avoiding the overlay of source and field
        index = np.where(abs(m-1)<1e-10)
    
        
        xx1 = x_plasma; # may modify, not exist outside [if ~isempty(index)]
        if index is not None:
            xx1[index] = x_plasma[index] + gapx;    # move 1 mm outside
            R1 = np.sqrt((xx1+x2)**2 + (y_plasma-y2)**2)
            m  = 4 * x2 * xx1 / R1**2;
        
        
        myk = special.ellipk(m)
        mye = special.ellipe(m)
        mi_plasma_limiter = self.cmu * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_plasma_limiter
        
        
        
    def calculate_mi_plasma_plasma(self, x_plasma, y_plasma, x2, y2, gapx):


        R1 = np.sqrt((x_plasma+x2)**2 + (y_plasma-y2)**2);
        m = 4 * x2 * x_plasma / R1**2;

        # avoiding the overlay of source and field
        index = np.where(abs(m-1)<1e-10)
        xx1 = x_plasma; # may modify, not exist outside [if ~isempty(index)]
        if index is not None:
            xx1[index] = x_plasma[index] + gapx;    # move 1 mm outside
            R1 = np.sqrt((xx1+x2)**2 + (y_plasma-y2)**2)
            m  = 4 * x2 * xx1 / R1**2;
        
        myk = special.ellipk(m)
        mye = special.ellipe(m)
        mi_plasma_plasma = self.cmu * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_plasma_plasma
    

            
        
    def calculate_mi_plasma_pf(self, x_plasma, y_plasma, index, gapx):

        
        # get pf location
        x_pf, y_pf, A_turn_coil = self.get_pf_location(index)
        
        
        mi_plasma_pf = 0*x_plasma
        for i in range(len(x_pf)):
            
            R1 = np.sqrt((x_plasma+x_pf[i])**2 + (y_plasma-y_pf[i])**2);
            m = 4 * x_pf[i] * x_plasma / R1**2;
    
            # avoiding the overlay of source and field
            index = np.where(abs(m-1)<1e-10)
            xx1 = x_plasma; # may modify, not exist outside [if ~isempty(index)]
            if index is not None:
                xx1[index] = x_plasma[index] + gapx;    # move 1 mm outside
                R1 = np.sqrt((xx1+x_pf[i])**2 + (y_plasma-y_pf[i])**2)
                m  = 4 * x_pf[i] * xx1 / R1**2;
            
            myk = special.ellipk(m)
            mye = special.ellipe(m)
            mi_plasma_pf = mi_plasma_pf + A_turn_coil[i] * self.cmu * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_plasma_pf
    
    
    def calculate_mi_limiter_pf(self, x_limiter, y_limiter, index, gapx):

        
        # get pf location
        x_pf, y_pf, A_turn_coil = self.get_pf_location(index)
        
        
        mi_limiter_pf = 0*x_limiter
        for i in range(len(x_pf)):
            
            R1 = np.sqrt(((x_limiter+x_pf[i])**2 + (y_limiter-y_pf[i])**2).astype('float'));
            m = 4 * x_pf[i] * x_limiter / R1**2;
    
            # avoiding the overlay of source and field
            index = np.where(abs(m-1)<1e-10)
        
            
            xx1 = x_limiter; # may modify, not exist outside [if ~isempty(index)]
            if index is not None:
                xx1[index] = x_limiter[index] + gapx;    # move 1 mm outside
                R1 = np.sqrt(((xx1+x_pf[i])**2 + (y_limiter-y_pf[i])**2).astype('float'))
                m  = 4 * x_pf[i] * xx1 / R1**2;
            
            
            myk = special.ellipk(m.astype('float'))
            mye = special.ellipe(m.astype('float'))
            mi_limiter_pf = mi_limiter_pf + A_turn_coil[i] * self.cmu * R1 * (2 * (myk-mye)-m * myk);
        
        return mi_limiter_pf
    
    
    def get_pf_location(self, index):
        
        """
        get pf grid location for green funciton calculation
        
        
        """
        
        # x_center = self.machine['PF'].[]
        
        PF = self.machine['PF'].iloc[index]
        
        x_center = PF['R/m']
        y_center = PF['Z/m']
        W = PF['W/m']
        H = PF['H/m']
        Nx = PF['Nx']
        Ny = PF['Ny']
        N = PF['N（匝数）']
        angle = PF['ANGLE']
        
        Nxy = Nx*Ny
        
        A_turn_coil = np.ones(Nxy) * N / (Nxy)
        step_x = W/Nx/2
        step_y = H/Ny/2
        
        x = np.arange(-W/2+step_x, W/2, 2*step_x)
        y = np.arange(-H/2+step_y, H/2, 2*step_y)
        [x, y] = np.meshgrid(x, y)
        
        
        
        if y_center>0:
            x = x - y/np.tan(angle*np.pi/180)
        elif y_center<=0:
            x = x + y/np.tan(angle*np.pi/180)
            
        x = x + x_center
        y = y + y_center
        
        x_pf = np.reshape(x, newshape=(Nxy))
        y_pf = np.reshape(y, newshape=(Nxy))
        A_turn_coil = np.reshape(A_turn_coil, newshape=(Nxy))
            
        
        return x_pf, y_pf, A_turn_coil
    
    
    
    
if __name__ == '__main__':

    
    GF = green_function()
    
    machine = GF.get_machine_parameters()
    
    eq = GF.update_green_function()
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    